const std = @import("std");

const Ast = @import("Ast");
const os = @import("os");
const fs = os.fs;
const MAXARG = os.MAXARG;
const PIPESIZE = fs.Pipe.SIZE;

const sys = @import("lib/syscall.zig");
const lib = @import("lib/user.zig");
const print = lib.print;
const color = lib.color;
const fd_t = sys.fd_t;
const c_str = sys.c_str;
pub const panic = lib.panic;

const String = []const u8;
const FdList = std.ArrayList(fd_t);
const RedirList = struct { pipe: [2]fd_t, fds: FdList = .empty };

const CmdError = error{OutOfMemory};
pub const Builtin = enum { cd, exit, echo, pwd, command };

var arena_allocator: std.heap.ArenaAllocator = .init(lib.heap.raw_c_allocator);
const allocator = arena_allocator.allocator();
const prompt = "$ ";
var cmd_buffer: [1024:0]u8 = undefined;

const usage =
    \\Usage: shell [OPTION]... [FILE]...
    \\
    \\Options:
    \\  --help         Display this help message and exit
    \\
    \\Execute commands from:
    \\  (no args)      Start interactive REPL shell
    \\  FILE           Run commands from each specified file
;

export fn main(argc: usize, argv: [*]const c_str) noreturn {
    errdefer |err| {
        print("sh: {t}\n", .{err});
        sys.exit(1);
    }

    if (argc > 1) {
        if (std.mem.eql(u8, std.mem.sliceTo(argv[1], 0), "--help")) {
            print("{s}\n", .{usage});
        } else {
            for (argv[1..argc]) |file| {
                try script(file);
            }
        }
        sys.exit(0);
    }

    try repl();
}

/// Read-Eval-Print Loop for interactive shell mode.
/// Loads .shrc if present, then continuously prompts for and executes commands.
fn repl() !noreturn {
    // attempts to load .shrc before entering repl
    var status: i32 = 0;
    if (sys.access(".shrc")) {
        script(".shrc") catch |err| switch (err) {
            error.CmdFailed => {
                print("failed to run .shrc script\n", .{});
                status = 1;
            },
            else => return err,
        };
    }

    while (true) {
        _ = arena_allocator.reset(.retain_capacity);
        const prompt_color = if (status == 0) color.yellow else color.red;
        print("{f}{s}{f}", .{ prompt_color, prompt, color.reset });

        var stdin: lib.File.Reader = .init(0, &cmd_buffer);
        const line = stdin.readline() orelse sys.exit(0);
        cmd_buffer[line.len] = 0; // chop \n
        const cmd: [:0]const u8 = cmd_buffer[0..line.len :0];

        status = try runcmd(cmd, true);
    }
}

/// Execute commands from a script file.
/// Opens the specified file and runs each line as a command sequentially.
fn script(path: c_str) !void {
    const n = sys.open(path, .{ .access = .RDONLY });
    if (n < 0) return error.FileNotFound;

    const fd: fd_t = @intCast(n);
    defer _ = sys.close(fd);

    var file: lib.File.Reader = .init(fd, &cmd_buffer);
    while (file.readline()) |line| {
        _ = arena_allocator.reset(.retain_capacity);
        cmd_buffer[line.len] = 0;
        const cmd: [:0]const u8 = cmd_buffer[0..line.len :0];
        if (try runcmd(cmd, false) != 0) return error.CmdFailed;
    }
}

/// Parse and execute a single command string.
fn runcmd(cmd: [:0]const u8, debug: bool) !i32 {
    // parse command
    var error_token: String = undefined;
    var tree = Ast.parse(allocator, cmd, &error_token) catch |err| switch (err) {
        error.EmptyCmd => return 0,
        error.TokenizeError, error.ParseError => {
            if (debug) {
                const position = @intFromPtr(error_token.ptr) - @intFromPtr(cmd.ptr);
                print("{s:[2]}{s:~<[3]}\n", .{ "", "^", position + prompt.len, error_token.len });
                print("sh: parse error near `{s}'\n", .{error_token});
            }
            return 1;
        },
        else => return err,
    };
    defer tree.deinit(allocator);

    // run command
    return runnode(tree, .root, true);
}

/// Recursively execute an AST node representing a command.
fn runnode(tree: Ast, index: Ast.Node.Index, root: bool) CmdError!i32 {
    var status: i32 = 0;
    const tag = tree.nodeTag(index);
    const data = tree.nodeData(index);

    // Process creation strategy:
    // - Root-level commands may need child processes (except builtins/lists)
    // - Nested commands always execute in current process
    const state: enum { root, parent, child } = if (root)
        switch (tag) {
            .list, .builtin, .@"and", .@"or" => .root,
            else => if (fork() == 0) .child else .parent,
        }
    else
        .child;

    // Parent process cleanup: wait for any forked children
    if (state == .parent) {
        _ = sys.wait(&status);
        return status;
    }

    // Command execution logic (runs in either root shell or child process)
    switch (tag) {
        .list => {
            const cmds = tree.extraDataSlice(data.extra_range, Ast.Node.Index);
            for (cmds) |cmd| {
                status = try runnode(tree, cmd, state == .root);
            }
        },
        .@"and", .@"or" => {
            status = try runnode(tree, data.node_and_node[0], state == .root);
            if (tag == .@"and" and status != 0) return status;
            if (tag == .@"or" and status == 0) return status;
            status = try runnode(tree, data.node_and_node[1], state == .root);
        },
        .builtin, .exec => {
            const tokens = tree.tokens.items(.lexeme)[data.token_range.start..data.token_range.end];
            if (tag == .builtin) {
                status = try builtin(std.meta.stringToEnum(Builtin, tokens[0]).?, tokens[1..]);
            } else {
                execute(tokens);
            }
        },
        .back => if (fork() == 0) {
            _ = try runnode(tree, data.node, true);
        },
        .pipe => try pipe(tree, data.node_and_node),
        .redir, .piped_redir => try redirect(tree, tag, data.node_and_extra),
        else => unreachable,
    }

    if (state == .child) sys.exit(0);
    return status;
}

/// Execute builtin shell commands.
/// These execute in the current shell process without forking.
fn builtin(tag: Builtin, tokens: []const String) error{OutOfMemory}!i32 {
    var status: i32 = 0;
    switch (tag) {
        .cd => {
            const path = try allocator.dupeZ(u8, tokens[0]);
            defer allocator.free(path);
            if (sys.chdir(path) < 0) {
                print("sh: cannot cd {s}\n", .{path});
                status = 1;
            }
        },
        .exit => {
            const n = if (tokens.len == 0) 0 else std.fmt.parseInt(i32, tokens[0], 0) catch 1;
            sys.exit(n);
        },
        .echo => {
            for (tokens) |token| {
                print("{s} ", .{token});
            }
            print("\n", .{});
        },
        .pwd => {
            var buf: [os.fs.MAXPATH]u8 = undefined;
            _ = sys.getcwd(&buf, buf.len);
            print("{s}\n", .{std.mem.sliceTo(&buf, 0)});
        },
        .command => {
            if (fork() == 0) execute(tokens);
            _ = sys.wait(&status);
        },
    }
    return status;
}

/// Execute external program by forking and replacing process image.
fn execute(tokens: []const String) noreturn {
    defer sys.exit(1);
    errdefer |err| print("exec {s} failed: {t}\n", .{ tokens[0], err });

    if (tokens.len > os.MAXARG) {
        print("sh: too many args\n", .{});
        return;
    }

    var list = try allocator.allocSentinel(?c_str, tokens.len, null);
    defer allocator.free(list);

    for (tokens, 0..) |token, i| {
        list[i] = try allocator.dupeZ(u8, token);
    }

    try lib.execvp(list[0].?, list);
}

/// Execute a command pipeline (cmd1 | cmd2).
/// Creates pipe, forks processes for both commands, and connects stdin/stdout.
fn pipe(tree: Ast, nodes: [2]Ast.Node.Index) !noreturn {
    var status: i32 = undefined;
    var p: [2]fd_t = undefined;
    _ = sys.pipe(&p);

    for (nodes, [2]u8{ 1, 0 }) |cmd, fd| {
        if (fork() == 0) {
            _ = sys.close(fd);
            _ = sys.dup(p[fd]);
            _ = sys.close(p[0]);
            _ = sys.close(p[1]);
            _ = try runnode(tree, cmd, false);
        }
    }

    _ = sys.close(p[0]);
    _ = sys.close(p[1]);
    _ = sys.wait(null);
    _ = sys.wait(&status);
    sys.exit(status);
}

/// Handle command with I/O redirection (<, >, >>, 2>, 2>>).
/// Sets up file redirection and piping for input/output streams.
fn redirect(tree: Ast, tag: Ast.Node.Tag, node_and_extra: struct { Ast.Node.Index, Ast.ExtraIndex }) !noreturn {
    const redir = tree.extraData(node_and_extra[1], Ast.Node.Redirection);

    var redir_list = rl: {
        var list: [3]?RedirList = @splat(null);

        for (0..3, [3]bool{
            redir.stdin != .none,
            redir.stdout != .none or redir.stdout_append != .none or tag == .piped_redir,
            redir.stderr != .none or redir.stderr_append != .none,
        }) |i, exist| {
            if (exist) {
                list[i] = .{
                    .pipe = pipe: {
                        var p: [2]fd_t = undefined;
                        _ = sys.pipe(&p);
                        break :pipe p;
                    },
                };
            }
        }

        if (tag == .piped_redir) try list[1].?.fds.append(allocator, 1);
        break :rl list;
    };

    // fork a new process to run the command
    if (fork() == 0) {
        for (redir_list, 0.., [3]usize{ 0, 1, 1 }) |list, stdio, i| {
            // redirect stdio to specified pipes and close original descriptors
            if (list) |ls| {
                _ = sys.close(@intCast(stdio));
                _ = sys.dup(ls.pipe[i]);
                _ = sys.close(ls.pipe[0]);
                _ = sys.close(ls.pipe[1]);
            }
        }
        _ = try runnode(tree, node_and_extra[0], false);
    }

    // fork two shells to process stdin/stdout redirections
    inline for (.{ .stdin, .stdout }) |stdio| {
        if (fork() == 0) {
            try redirectIO(tree, redir, &redir_list, stdio);
            sys.exit(0);
        }
    }

    // the parent shell processes stderr redirections
    try redirectIO(tree, redir, &redir_list, .stderr);

    var status: i32 = 0;
    for (0..3) |_| {
        var st: i32 = undefined;
        _ = sys.wait(&st);
        status += st;
    }
    sys.exit(status);
}

/// Process I/O redirection between shell and executed command.
/// Manages data piping between file descriptors for stdin/stdout/stderr redirection.
fn redirectIO(
    tree: Ast,
    redir: Ast.Node.Redirection,
    list: *[3]?RedirList,
    comptime stdio: enum(u8) { stdin, stdout, stderr },
) !void {
    const idx: usize = @intFromEnum(stdio);

    // close unused pipes
    for (list, 0..) |redir_list, i| {
        if (idx == i) continue;

        if (redir_list) |ls| {
            _ = sys.close(ls.pipe[0]);
            _ = sys.close(ls.pipe[1]);
        }
    }

    // open input/output files
    inline for (
        std.meta.fields(Ast.Node.Redirection),
        [5]fs.O{
            .{ .access = .RDONLY }, // <
            .{ .access = .WRONLY, .create = true, .truncate = true }, // >
            .{ .access = .WRONLY, .create = true, .append = true }, // >>
            .{ .access = .WRONLY, .create = true, .truncate = true }, // 2>
            .{ .access = .WRONLY, .create = true, .append = true }, // 2>
        },
        .{ 0, 1, 1, 2, 2 },
    ) |field, mode, i| {
        if (idx != i) continue;

        if (@field(redir, field.name).unwrap()) |node| {
            const files = tree.extraDataSlice(tree.nodeData(node).extra_range, Ast.TokenIndex);
            const tokens = tree.tokens.items(.lexeme);

            for (files) |file| {
                const path = try allocator.dupeZ(u8, tokens[file]);
                defer allocator.free(path);
                try list[idx].?.fds.append(allocator, @intCast(sys.open(path, mode)));
            }
        }
    }

    var pipe_buffer: [PIPESIZE]u8 = undefined;

    if (stdio == .stdin) {
        // read from input files and write to child stdin pipe
        if (list[0]) |ls| {
            _ = sys.close(ls.pipe[0]);
            defer _ = sys.close(ls.pipe[1]);

            for (ls.fds.items) |fd| {
                while (true) {
                    const nbytes = sys.read(fd, &pipe_buffer, PIPESIZE);
                    if (nbytes <= 0) break;
                    _ = sys.write(ls.pipe[1], &pipe_buffer, @intCast(nbytes));
                }
            }
        }
    } else for (list[1..], 1..) |redir_list, i| {
        // read from child stdout/stderr pipes and write to output files
        if (idx != i) continue;

        if (redir_list) |ls| {
            _ = sys.close(ls.pipe[1]);
            defer _ = sys.close(ls.pipe[0]);

            const fds = ls.fds.items;
            while (true) {
                const nbytes = sys.read(ls.pipe[0], &pipe_buffer, PIPESIZE);
                if (nbytes <= 0) break;

                for (fds) |fd| {
                    _ = sys.write(fd, &pipe_buffer, @intCast(nbytes));
                }
            }
        }
    }
}

fn fork() sys.pid_t {
    const pid = sys.fork();
    if (pid == -1) @panic("fork");
    return @intCast(pid);
}
